(ECON) Solving Systems in Julia, While Keeping a Tight Link to the Math

Introduction

The main conclusions of the post are:

Warning!

The aspects regarding Julia don't require any knowledge of Economics. The post basically is about solutions for a non-linear system of equations. If you're only interested in the Julia implementation and don't care about Economics, go directly to this section. There, I start by summarizing the system of equations to solve, regardless of how they were obtained.

Setup

There is a set of firms $\Omega$. The notation for a given firm $i$

• $q_{i}$ $\rightarrow$ $i$'s quantity

• $p_{i}$ $\rightarrow$ $i$'s price

• $z_{i}$ $\rightarrow$ $i$'s "quality" (exogenous).

Moreover, the industry values are:

• $E$ $\rightarrow$ industry expenditure (exogenous)

• $\mathbb{P}$ $\rightarrow$ price index

• $\mathbb{W}:=\frac{E}{\mathbb{P}}$ $\rightarrow$ welfare

• $S_{m}$ $\rightarrow$ total market share of imports

We'll make use of a CES demand function. This is well-suited for expressing all our variables in percentage variations through the "hat-algebra" procedure. Specifically, the demand for firm $i$ is

where the price index is given by $\mathbb{P}:=\left[\sum_{i\in\Omega}\left(z_{i}\right)^{\delta}\left(p_{i}\right)^{1-\sigma}\right]^{\frac{1}{1-\sigma}}$.

This demand determines a price elasticity given by

Given a revenue $r_{i}:=q_{i}p_{i}$, the market share of $i$ defined by $s_{i}:=\frac{r_{i}}{E_{j}}$. This can be expressed as $s_{i}=\left(z_{i}\right)^{\delta}\left(p_{i}\right)^{1-\sigma}\left(\mathbb{P}\right)^{\sigma-1}.$, which substituting in the definition of the price index becomes

Optimal Prices and Equilibrium

Optimal price of a firm $i$ can be expressed as a function of its market share by using that

where $\mu\left(\varepsilon_{i}\right)=\frac{\varepsilon_{i}}{\varepsilon_{i}-1}$ and $\varepsilon\left(s_{i}\right)=\sigma-s_{i}\left(\sigma-1\right)$.

Finally, there's equilibrium in the industry when the sum of market shares equals one:

Summary of Equations to be Solved

Identifying $\widehat{\mathbb{P}}$ can be done by solving for $\left(\widehat{\mathbb{P}},\left(\widehat{s}_{i}\right)_{i\in\Omega}\right)$ simultaneously. This requires solving the following system of equations,

where (2) holds for each $i$ and

Note that $\widehat{p}_{i}\left(\widehat{\varepsilon}_{i}\right)$ and $\widehat{\varepsilon}_{i}\left(\widehat{s}_{i}\right)$, so that they're fully determined by $\widehat{s}_{i}$.

Information at our disposal to solve the system:

• $s_{i}^{*}$ for each $i\in\Omega$.

• $\varepsilon_{i}^{*}$ for each $i\in\Omega$, computed by knowledge of $s_{i}^{*}$.

• Parameters: only $\sigma$, for which we'll assume a value.

• Shock: $\Delta S_{m}$, for which we'll assume a value.

Solving the System: For-Loops

The package NLsolve requires specifying the equations to be solved through residuals. Each residual represents the difference between the value provided by an equation against its value under a proposed solution. The code includes the residuals through a vector res, whose entries need to be zero in equilibrium. Likewise, the solutions are represented by a vector called sol. The rest of the variables in the code use the same notation as the equations, thanks to the use of Unicode characters.

One downside of NLsolve is that its solver function must depend only on the residuals and the solutions—you can't add parameters to the function. Therefore, to achieve type stability, you need to wrap the solver and computation functions within the same function.

We start with a baseline implementation relying on for-loops. This will help us explain what the code does.

function computation(sᵢ, ΔSₘ, eᵢ, σ)
    êᵢ = similar(sᵢ); p̂ᵢ = similar(sᵢ)

    function solver!(res, sol)
        P̂  = sol[1]
        ŝᵢ  = sol[2:length(sᵢ)+1]

        for i in eachindex(êᵢ)
            êᵢ[i] = (σ - (σ - 1) * sᵢ[i] * ŝᵢ[i]) / (σ - (σ - 1) * sᵢ[i])
            p̂ᵢ[i] = êᵢ[i] * (eᵢ[i] - 1) / (êᵢ[i] * eᵢ[i] - 1)
        end

        sum_shares = [(1 - ŝᵢ[i]) * sᵢ[i] for i in eachindex(sᵢ)]
        res[1]     = sum(sum_shares) - ΔSₘ

        for i in 2:length(sᵢ)+1
            res[i] = ŝᵢ[i-1] - p̂ᵢ[i-1]^(1 - σ) * P̂^(σ - 1)
        end

    end

    function execute()
        sol0     = vcat(0.9, 0.8 .* ones(length(sᵢ)))
        solution = nlsolve(solver!, sol0)

        converged(solution) || error("Failed to converge")
        return solution.zero
    end

    return execute()
end

  50.182 ms (8094 allocations: 41.60 MiB)

Next, we'll propose several ways to express the function solver! and its components, which capture the system of equations and variables to be solved. Most of the computational time is spent on the multiple calls to solver!. As the rest of the terms have a negligible impact on time, it's possible to reexpress them if that increases readability.

Moreover, the function execute sets the initial conditions and the call to the solver. It'll stay the same in all the code examples.

The first entry sol[1] represents $\widehat{\mathbb{P}}$, while the other entries in sol refer to $\widehat{s}_{i}$ for each $i\in\Omega$. Likewise, res[1] expresses (1), with the rest of the entries in res referring to equation (2) for each $i\in\Omega$.

Improvements

The improvements proposed are below.

function computation(sᵢ, ΔSₘ, eᵢ, σ)
    êᵢ, p̂ᵢ = (similar(sᵢ) for _ in 1:2)

    function solver!(res, sol)
        P̂  = sol[1]
        ŝᵢ  = @view sol[2:length(sᵢ)+1]

        sum_shares = ((1 - ŝᵢ[i]) * sᵢ[i] for i in eachindex(sᵢ))
        res[1]     = sum(sum_shares) - ΔSₘ

        rest  = @view res[2:length(sᵢ)+1]          # to update the results below
            @turbo for i in eachindex(êᵢ)
                êᵢ[i]     =  (σ - (σ - 1) * sᵢ[i] * ŝᵢ[i]) / (σ - (σ - 1) * sᵢ[i])
                p̂ᵢ[i]     =  êᵢ[i] * (eᵢ[i] - 1) / (êᵢ[i] * eᵢ[i] - 1)
                rest[i]  =  ŝᵢ[i] - p̂ᵢ[i]^(1 - σ) * P̂^(σ - 1)
            end
    end

    function execute()
        sol0     = vcat(0.9, 0.8 .* ones(length(sᵢ)))
        solution = nlsolve(solver!, sol0)

        converged(solution) || error("Failed to converge")
        return solution.zero
    end

    return execute()
end

  29.187 ms (64 allocations: 9.74 MiB)

The modifications and their justifications are the following. First, the line êᵢ = similar(sᵢ); p̂ᵢ = similar(sᵢ) is wordy, and can be simplified by using êᵢ, p̂ᵢ = (similar(sᵢ) for _ in 1:2). This makes it possible to initialize multiple variables in one line. Note we could've also used an array comprehension instead of a generator, with a minimal impact on performance (the comprehension is only called once as it's not part of solver!).

Second, slices consisting of a scalar don't allocate memory, while slices with more than two elements do. This determines that the slice of sol comprising each $\widehat{s}_{i}$ allocates. We circumvent this by the use of @view.

Third, we're not interested in sum_shares itself. Rather, it acts as an auxiliary that is eventually called to compute sum(sum_shares). In the baseline case, sum_shares is defined through an array comprehension. This stores the result in a new vector, and therefore allocates memory. Instead, we can delay its computation through a generator, allowing the result to fuse with sum—Julia can avoid the allocation when it knows that the vector will ultimately be "reduced" to a scalar.

Fourth, the package LoopVectorization provides the macro @turbo. This significantly improves the speed of for-loops and vectorized code when the operations involve multiple terms.

Fifth, all the for-loops to compute the entries greater than 2 of res can be combined into a single for-loop. This requires defining a view of the object, which we call rest, enabling us to include rest and the auxiliaries for calculating $\widehat{\varepsilon}_{i}$ and $\widehat{p}_{i}$.

Combining these computations into one for-loop is crucial for efficiently parallelizing the code, as it gathers the three most demanding calculations of solver!. The number of firms in our example is too small to see any gain from parallelizing code. Nonetheless, it's worth remarking that, if necessary, we can combine the advantages of both parallelizing code and applying @turbo through the macro @tturbo.

Vectorized Code: My Preferred Solution

We conclude by providing an alternative for the last code presented. It improves readability further by separating the system of equations to be solved and the definition of the intermediate variables, $\widehat{p}_{i}$ and $\widehat{\varepsilon}_{i}$. The code is also well-suited for cases where the computations of $\widehat{p}_{i}$ and $\widehat{\varepsilon}_{i}$ are more complex.

The approach relies on in-place functions for $\widehat{p}_{i}$ and $\widehat{\varepsilon}_{i}$. The function solver! uses them by initializing $\widehat{p}_{i}$ and $\widehat{\varepsilon}_{i}$, and updating these variables in each iteration.

The auxiliary functions are defined below. We provide two alternatives. The second one in particular simply shows that we could repeat the vectorized form from the previous code.

function êᵢ!(êᵢ, ŝᵢ, sᵢ, σ)
    @turbo for i in eachindex(êᵢ)
        numerator   = σ - (σ - 1) * sᵢ[i] * ŝᵢ[i]
        denominator = σ - (σ - 1) * sᵢ[i]
        
        êᵢ[i] = numerator / denominator
    end
end

function p̂ᵢ!(p̂ᵢ, êᵢ, eᵢ)
    @turbo for i in eachindex(êᵢ)
        numerator   = êᵢ[i] * (eᵢ[i] - 1)
        denominator = êᵢ[i] * eᵢ[i] - 1

        p̂ᵢ[i] = numerator / denominator
    end   
end

êᵢ!(êᵢ, ŝᵢ, sᵢ, σ) = @turbo @. êᵢ = (σ - (σ - 1) * sᵢ * ŝᵢ) / (σ - (σ - 1) * sᵢ)
p̂ᵢ!(p̂ᵢ, êᵢ, eᵢ)    = @turbo @. p̂ᵢ = êᵢ * (eᵢ - 1) / (êᵢ * eᵢ - 1)

The solver then initializes êᵢ and p̂ᵢ, and updates them in each iteration by executing êᵢ! and p̂ᵢ!. The rest of the code is identical to the previous section.

function computation(sᵢ, ΔSₘ, eᵢ, σ)
    êᵢ, p̂ᵢ      = (similar(sᵢ) for _ in 1:2)
    idx_shares = 2:length(sᵢ)+1

    @views function solver!(res, sol)
        P̂  = sol[1]
        ŝᵢ  = sol[idx_shares]
        
        êᵢ!(êᵢ, ŝᵢ, sᵢ, σ)
        p̂ᵢ!(p̂ᵢ, êᵢ, eᵢ)        

                  res[1]          = sum(@~ (1 .- ŝᵢ) .* sᵢ) - ΔSₘ
        @turbo @. res[idx_shares] = ŝᵢ - p̂ᵢ^(1 - σ) * P̂^(σ - 1)
    end

    function execute()
        sol0     = vcat(0.9, 0.8 .* ones(length(sᵢ)))
        solution = nlsolve(solver!, sol0)

        converged(solution) || error("Failed to converge")
        return solution.zero
    end

    return execute()
end

  28.925 ms (67 allocations: 9.74 MiB)

Note the use of @views before the function statement. This is a neat way to indicate that each slice within the function should be treated as a view, rather than a copy.